#include "mol_sys.h"
#include "physical_properties.h"
#include "autocorr_CUDA.cuh"
#include "force.h"
#include <iostream> 

/**************************************************************************Classic**************************************************************************/
mol_sys::mol_sys() {
    Natoms = Parameter_Natoms;
    temperature = Parameter_Temperature_imposed;
    potential_energy = 0;
    kenetic_energy = 0;
    total_energy = 0;
    for (int i = 0; i < Natoms; i++) {
        Atom at;
        mol.push_back(at);
    }
}

double mol_sys::Compute_sys_Ek(const std::vector<Atom> mol) {
    double Ek = 0;
    for (int i = 0; i < Parameter_Natoms; i++) {
        double E = 0;
        for (int j = 0; j < 3; j++) {
            E += mol[i].veloc[j] * mol[i].veloc[j];
        }
        E *= 0.5 * mol[i].mass;                             // Here mass should be in AU
        Ek += E;
    }
    
    return Ek;
}

double mol_sys::Compute_sys_Ep(const std::vector<Atom>mol) {
    double potential_energy = 0.0;
    int Natoms = mol.size();
    double dx = 0, dy = 0, dz = 0, r = 0, r2 = 0, attractive_term = 0, repulsive_term = 0;

    for (int i = 0; i < Natoms - 1; i++) {
        for (int j = i + 1; j < Natoms; j++) {
            LJ_Force::distance(mol[i].coord, mol[j].coord, dx, dy, dz, r);
            if (r2 < Parameter_Rcutoff) {
                LJ_Force::terms(r, attractive_term, repulsive_term);
                double energyij = LJ_Force::energy(attractive_term, repulsive_term);
                potential_energy = potential_energy + energyij;
            }
        }
    }

    return potential_energy;

}

double mol_sys::Compute_Temperature(const std::vector<Atom>& mol, const double Ek) {
    double temperature = 2.0 * Ek / (3.0 * Constant_Boltz * Parameter_Natoms); // temperature in K
    return temperature;
}

void mol_sys::Compute_Temperature_and_Pressure(const std::vector<Atom>& mol, double& Ek, double& temperature, double& configurational_pressure) {
    temperature = 2.0 * Ek / (3.0 * Constant_Boltz * Parameter_Natoms);         // temperature in K
    configurational_pressure = 2.0 * Ek / (3.0 * pow(Parameter_Box_size, 3)); 

}

double mol_sys::compute_velocities_average(std::vector<Atom> mol) {
    double velocities_average = 0;
    for (int i = 0; i < Parameter_Natoms; i++)
        velocities_average += sqrt(pow(mol[i].veloc[0], 2) + pow(mol[i].veloc[1], 2) + pow(mol[i].veloc[2], 2));
    velocities_average /= Parameter_Natoms;
    return velocities_average;
}

double mol_sys::compute_forces_average(std::vector<Atom> mol) {
    double froces_average = 0;
    for (int i = 0; i < Parameter_Natoms; i++)
        froces_average += sqrt(pow(mol[i].force[0], 2) + pow(mol[i].force[1], 2) + pow(mol[i].force[2], 2));
    froces_average /= Parameter_Natoms;
    return froces_average;
}

std::vector<double> mol_sys::compute_coord_average(std::vector<Atom> mol) {
    std::vector<double>coord_average(3, 0);
    for (int i = 0; i < Parameter_Natoms; i++) {
        for (int j = 0; j < 3; j++)
            coord_average[j] += mol[i].coord[j];
    }
    for (int i = 0; i < 3; i++)
        coord_average[i] /= Parameter_Natoms;
    return coord_average;
}

double mol_sys::autocorr(const int Tau, const int seq) {
    double acf = 0;
    std::vector<double> elem(Parameter_Nsteps);

    for (int step = 0; step < Parameter_Nsteps; step++)
        elem[step] = velocitiesOFeach[step].mol[seq].veloc[0] + velocitiesOFeach[step].mol[seq].veloc[1] + velocitiesOFeach[step].mol[seq].veloc[2];

    for (int low = 0, high = Tau; high < Parameter_Nsteps; low++, high++) 
        acf += elem[low] * elem[high];

     acf /= Parameter_Nsteps - Tau;
    return acf;
}

double mol_sys::autocorr_4_sys(const int tau) {
    double acf = 0;

    for (int i = 0; i < Parameter_Natoms; i++)
        //acf += autocorr_CUDA(tau, i);
        acf += mol_sys::autocorr(tau, i);

    acf /= Parameter_Natoms;

    return acf;
}

 std::vector<double> mol_sys::autocorr_store() {
    std::vector<double> acf(Parameter_Nsteps, 0);

    for (int tau = 0; tau < Parameter_Nsteps; tau++) {
        acf[tau] = mol_sys::autocorr_4_sys(tau);
    }

    //double value = acf[0];

    //for (int tau = 0; tau < Parameter_Nsteps; tau++) {
    //    acf[tau] /= value;
    //} 

    return acf;
}

double mol_sys::Diffusion_Coefficient(vector<double> acf) {
    double diff = 0;

    for (int tau = 0; tau < 400; tau++)
        diff += acf[tau];

   /* for (int i = 0; i < acf.size(); i++)
        diff += acf[i] * Parameter_Timestep_au;*/

    //diff /= 3;

    //diff *= pow(Unit_length_au2m, 2) / (Unit_time_au2fs * 1E15);

    //diff /= 8637.973149;

    return diff;
}
//double mol_sys::Diffusion_Coefficient() {
//    double diff = 0;
//
//    std::vector<double>acf(Parameter_Nsteps, 0);
//    std::vector<double>acf_temp(Parameter_Nsteps, 0);
//
//    for (int step = 0; step < Parameter_Nsteps; step++) {
//        acf_temp = autocorr(Parameter_Nsteps);						//use a temporary array to store the result of autocorrelation who collects values from one spesific atom at one moment
//        for (int k = 0; k < Parameter_Nsteps; k++) {
//            acf[k] += acf_temp[k];
//        }  
//    }
//
//    for (int i = 0; i < Parameter_Nsteps; i++) {
//        acf[i] /= Parameter_Natoms;
//    }
//
//    ofstream file;
//    file.open("..\\result\\acf_results.csv");
//    for (int i = 0; i < Parameter_Nsteps; i++) {
//        file << acf[i] << endl;
//    }
//
//    for (int i = 0; i < acf.size(); i++)
//        diff += acf[i] * Parameter_Timestep_au;
//
//    return diff;
//}





/**************************************************************************PIMD**************************************************************************/


/*************************energy calculating*************************/
double mol_sys::Compute_E_spr(std::vector< std::vector<Atom> > &mirr) {
    double E_spr = 0;
    for (int j = 0; j < 3; j++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            int m = 0;
            for (int k = 0; k < Parameter_nbeads; k++) {
                m = k + 1;
                if (k == Parameter_nbeads - 1)m = 0;
                E_spr += 0.5 * mirr[k][i].mass * pow(Constant_wp, 2) * pow((mirr[k][i].coord[j] - mirr[m][i].coord[j]), 2);
            }
        }
    }
    return E_spr;
}

double mol_sys::Compute_Ek_vir(std::vector< std::vector<Atom> > &mirr) {
    double Ek_vir = 0;
    for (int j = 0; j < 3; j++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int k = 0; k < Parameter_nbeads; k++) {
                Ek_vir += mirr[k][i].coord[j] * (-mirr[k][i].force[j]);
            }
        }
    }
    Ek_vir /= 2.0 * Parameter_nbeads;
    return Ek_vir;
}

double mol_sys::Compute_Ek_vircen(std::vector< std::vector<Atom> >&mirr) {
    vector< vector<double> > Mirror_Center = LJ_Force::Compute_Mirror_Center(mirror);
    double Ek_vircen = 0;
    for (int j = 0; j < 3; j++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int k = 0; k < Parameter_nbeads; k++) {
                Ek_vircen += (mirr[k][i].coord[j] - Mirror_Center[i][j]) * (-mirr[k][i].force[j]);
            }
        }
    }
    Ek_vircen /= 2.0 * Parameter_nbeads;
    return Ek_vircen;

}

double mol_sys::Compute_Ek_prim(std::vector< std::vector<Atom> > &mirr) {
    double Ek_prim = 3.0 * Parameter_Natoms * Parameter_nbeads / (2 * Constant_beta) - Compute_E_spr(mirror);
    return Ek_prim;
}

double mol_sys::Compute_Ek_stage(std::vector< std::vector<Atom> > &mirr) {
    double Ek_stage = 0;
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++)
            Ek_stage += mirr[k][i].mass * (pow(mirr[k][i].veloc[0], 2) + pow(mirr[k][i].veloc[1], 2) + pow(mirr[k][i].veloc[2], 2)) / 2;
    }
    return Ek_stage;
}

double mol_sys::Compute_E_spr_stage(std::vector< std::vector<Atom> > &mirr) {
    double E_spr_stage = 0;
    for (int k = 1; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++)
            E_spr_stage += mirr[k][i].mass * Constant_wp * Constant_wp * (pow(mirr[k][i].coord[0], 2) + pow(mirr[k][i].coord[1], 2) + pow(mirr[k][i].coord[2], 2)) / 2;
    }
    return E_spr_stage;
}

double mol_sys::Compute_Ep_stage(std::vector< std::vector<Atom> >& mirr) {
    double potential = 0, v = 0, Ep = 0;

    for (int k = 0; k < Parameter_nbeads; k++) {
        potential = mol_sys::Compute_sys_Ep(mirr[k]);

        Ep += potential / Parameter_nbeads;
    }
    return Ep;
}

//double mol_sys::hamiltonian_compute_staging() {
//    double  ham = 0, p1 = 0, p2 = 0, p3 = 0;
//
//    //p1 = Ek_stage
//    for (int k = 0; k < Parameter_nbeads; k++) {
//        for (int i = 0; i < Parameter_Natoms; i++)
//            p1 += staging_mirror[k][i].mass * (pow(staging_mirror[k][i].veloc[0], 2) + pow(staging_mirror[k][i].veloc[1], 2) + pow(staging_mirror[k][i].veloc[2], 2)) / 2;
//    }
//    //p1 *= Unit_energy_au2kJpermol;
//
//    //p2 = E_spr
//    for (int k = 1; k < Parameter_nbeads; k++) {
//        for (int i = 0; i < Parameter_Natoms; i++)
//            p2 += staging_mirror[k][i].mass * Constant_wp * Constant_wp * (pow(staging_mirror[k][i].coord[0], 2) + pow(staging_mirror[k][i].coord[1], 2) + pow(staging_mirror[k][i].coord[2], 2)) / 2;
//    }
//    //p2 *= Energy_kJpermol_reverse;
//    //p2 = 1000 * p2 * Energy_J_reverse * Avogadro;
//    //p3 = E_p
//
//    double potential = 0, v = 0;
//
//    //p3 = Ep_stage
//    for (int k = 0; k < Parameter_nbeads; k++) {
//        LJ_Force::Compute_Forces(mirror[k], potential, v);
//
//        p3 += potential / Parameter_nbeads;
//    }
//
//  /*  cout << "Ek" << '\t' << p1 << '\t'
//        << "Espr" << '\t' << p2 << '\t'
//        << "Ekprim" << '\t' << 1.5 * Parameter_Natoms * Parameter_nbeads / Constant_beta - p2 << '\t'
//        << "Ep" << '\t' << p3 << '\t'
//        << "Etot" << '\t' << p1 + p2 + p3 << endl;*/
//
//    double coord = staging_mirror[3][6].coord[1];
//    double velocity = staging_mirror[3][6].veloc[1];
//    
//    double t1 = staging_mirror[0][0].veloc[0];
//    double seq_bead = 0, seq_atom = 0;
//
//    for (int k = 0; k < Parameter_nbeads; k++) {
//        for (int i = 0; i < Parameter_Natoms; i++) {
//            if (t1 < staging_mirror[k][i].veloc[0]) {
//                t1 = staging_mirror[k][i].veloc[0];
//                seq_bead = k, seq_atom = i;
//            }
//        }
//    }
//
//
//    double t2 = (pow(staging_mirror[3][6].coord[0], 2) + pow(staging_mirror[3][6].coord[1], 2) + pow(staging_mirror[3][6].coord[2], 2)) / 2;
//    double t3 = potential;
//
//    cout << p1 << '\t' << p2 << '\t' << p3 << endl;
//
//    p3 *= Unit_energy_au2kJpermol;
//    ham = p1 + p2 + p3;
//
//    return ham;
//}

double mol_sys::hamiltonian_compute_staging() {
    double  ham = 0, p1 = 0, p2 = 0, p3 = 0;

    //p1 = Ek_stage
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++)
            p1 += staging_mirror[k][i].mass * (pow(staging_mirror[k][i].veloc[0], 2) + pow(staging_mirror[k][i].veloc[1], 2) + pow(staging_mirror[k][i].veloc[2], 2)) / 2;
    }
    


    //p2 = E_spr
    for (int k = 1; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++)
            p2 += staging_mirror[k][i].mass * Constant_wp * Constant_wp * (pow(staging_mirror[k][i].coord[0], 2) + pow(staging_mirror[k][i].coord[1], 2) + pow(staging_mirror[k][i].coord[2], 2)) / 2;
    }



    //p3 = Ep_stage
    double potential = 0, v = 0;

    for (int k = 0; k < Parameter_nbeads; k++) {
        //LJ_Force::Compute_Forces(mirror[k], potential, v);
        potential = mol_sys::Compute_sys_Ep(mirror[k]);

        p3 += potential / Parameter_nbeads;
    }

    //cout << p1 << '\t' << p2 << '\t' << p3 << endl;

    ham = p1 / Parameter_nbeads + p2 / Parameter_nbeads + p3;
    ham *= Unit_energy_au2kJpermol;

    //cout << ham << endl;

    return ham;
}




/*************************staging transform*************************/
void mol_sys::mass_staging_transform(std::vector<Atom> mol, std::vector< std::vector<Atom> >& stag_mirr) {
    for (int i = 0; i < Parameter_Natoms; i++) {
        stag_mirr[0][i].mass = mol[i].mass;
    }

    for (int k = 1; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++)
            stag_mirr[k][i].mass = (k + 1.0) / (double)k * mol[i].mass;
    }
}

void mol_sys::coord_staging_transform() {
    /*int m = 0;
    if (k == Parameter_nbeads - 1)m = 0;
    else m = k + 1;

    mol[0].coord[0] = mirror[k][0].coord[0];
    mol[0].coord[0] = mirror[k][1].coord[1];
    mol[0].coord[0] = mirror[k][2].coord[2];

    for (int i = 1; i < mol.size(); i++) {
        mol[i].coord[0] = mirror[k][i].coord[0] - (1 - 1.0 / (k + 1)) * mirror[m][i].coord[0] - (1.0 / (k + 1)) * mirror[0][i].coord[0];
        mol[i].coord[1] = mirror[k][i].coord[1] - (1 - 1.0 / (k + 1)) * mirror[m][i].coord[1] - (1.0 / (k + 1)) * mirror[0][i].coord[1];
        mol[i].coord[2] = mirror[k][i].coord[2] - (1 - 1.0 / (k + 1)) * mirror[m][i].coord[2] - (1.0 / (k + 1)) * mirror[0][i].coord[2];
    }*/
    for (int i = 0; i < Parameter_Natoms; i++) {
        for (int j = 0; j < 3; j++)
            staging_mirror[0][i].coord[j] = mirror[0][i].coord[j];
    }

    if (Parameter_nbeads > 1) {
        for (int k = 1, m = k + 1; k < Parameter_nbeads; k++) {                         //m = k+1, loop back to the head while k = nbead
            if (k == Parameter_nbeads - 1)m = 0;
            for (int i = 0; i < Parameter_Natoms; i++) {
                for (int j = 0; j < 3; j++) {
                    staging_mirror[k][i].coord[j] = mirror[k][i].coord[j] - (double)k / (k + 1.0) * mirror[m][i].coord[j] - mirror[0][i].coord[j] / (k + 1.0);
                    //staging_mirror[k][i].coord[j] = mirror[k][i].coord[j] - -(double)k / (k + 1.0) * mirror[m][i].coord[j];
                }
            }
        }
    }
}

void mol_sys::coord_inverse() {
    for (int i = 0; i < Parameter_Natoms; i++) {
        for (int j = 0; j < 3; j++)
            mirror[0][i].coord[j] = staging_mirror[0][i].coord[j];
    }
    if (Parameter_nbeads > 1) {
        for (int k = 1; k < Parameter_nbeads - 1; k++) {
            for (int i = 0; i < Parameter_Natoms; i++) {
                for (int j = 0; j < 3; j++)
                    mirror[k][i].coord[j] = staging_mirror[k][i].coord[j] + k / (k + 1.0) * mirror[k + 1][i].coord[j] + 1 / (k + 1.0) * staging_mirror[0][i].coord[j];
            }
        }
        int k = Parameter_nbeads - 1;
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++)
                mirror[k][i].coord[j] = staging_mirror[k][i].coord[j] + k / (k + 1.0) * mirror[0][i].coord[j] + 1 / (k + 1.0) * staging_mirror[0][i].coord[j];
        }
    }
}


void mol_sys::forces_staging_transform() {
    vector< vector<double> >temp(Parameter_Natoms, vector<double>(3, 0));
    for (int l = 0; l < Parameter_nbeads; l++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++)
                temp[i][j] += mirror[l][i].force[j];
        }
    }
    for (int i = 0; i < Parameter_Natoms; i++) {
        for (int j = 0; j < 3; j++) {
            staging_mirror[0][i].force[j] = temp[i][j] / Parameter_nbeads;
        }
    }

    for (int k = 1; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++)
                staging_mirror[k][i].force[j] = mirror[k][i].force[j] / Parameter_nbeads + (k - 1.0) / (double)k * staging_mirror[k - 1][i].force[j] ;
               // staging_mirror[0][i].force[j] = 0;
        }
    }
}

void mol_sys::velocities_update_PIMD_staging() {
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++)
                mirror[k][i].veloc[j] = staging_mirror[k][i].veloc[j];
        }
    }
}

void mol_sys::staging_transform() {
    vector< Atom> mol(256);
    for (int i = 0; i < Parameter_Natoms; i++) { mol[i].mass = Parameter_Mass_Au; }
    mass_staging_transform(mol, staging_mirror);

    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++)
                staging_mirror[k][i].veloc[j] = mirror[k][i].veloc[j];
        }
    }
    coord_staging_transform();
    forces_staging_transform();
}

double mol_sys::max_veloc() {
    double t = staging_mirror[0][0].veloc[0];
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++) {
                if (t <= staging_mirror[k][i].veloc[j])t = staging_mirror[k][i].veloc[j];
            }
        }
    }
    return t;
}

double mol_sys::max_coord(vector<vector<Atom>>mirr) {
    double t = mirr[0][0].coord[0];
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++) {
                if (t <= mirr[k][i].coord[j])t = mirr[k][i].coord[j];
            }
        }
    }
    return t;
}

double mol_sys::max_force(vector<vector<Atom>>mirr) {
    double t = mirr[0][0].force [0];
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++) {
                if (t <= mirr[k][i].force[j])t = mirr[k][i].force[j];
            }
        }
    }
    return t;
}